import torch

checkpoint_path = "/NAS6/Members/linchenxi/projects/crop_recognition/checkpoints/henan_shandong_corn/test_4/saved_checkpoint_latest.pth"  # noqa:E501


saving_obj = torch.load(checkpoint_path)
conf_matrix_ori = saving_obj.metric_dict["confusion_matrix"]


conf_matrix = torch.sum(conf_matrix_ori, dim=0)
valid_idx = (
    torch.sum(conf_matrix, dim=1) != 0
)  # (True Positive + False Negative) cannot be 0
true_positive = torch.diag(conf_matrix)  # (C,)
false_positive = torch.sum(conf_matrix, 0) - true_positive  # (C,)
false_negative = torch.sum(conf_matrix, 1) - true_positive  # (C,)

true_positive = true_positive[valid_idx]
false_positive = false_positive[valid_idx]
false_negative = false_negative[valid_idx]

iou = true_positive / (true_positive + false_positive + false_negative)  # (C,)
miou = float(torch.mean(iou))
precision = true_positive / (true_positive + false_positive)
precision[torch.isnan(precision)] = 1
avg_precision = float(torch.mean(precision))
recall = true_positive / (true_positive + false_negative)
avg_recall = float(torch.mean(recall))
f1_score = 2 * precision * recall / (precision + recall)
avg_f1_score = float(torch.prod(f1_score))

print()
